import numpy as np
import spatialmath as sm

from ..biped_robot import BipedRobot, fkine


def wrap_to_pi(theta: float) -> float:
    if theta < -np.pi:
        theta += 2 * np.pi
    elif theta > np.pi:
        theta -= 2 * np.pi
    return theta


class Cassie(BipedRobot):

    def __init__(self) -> None:
        super().__init__()

        # 左腿
        self.x1 = 0.021
        self.y1 = 0.1350
        self.z1 = 0.0

        self.x2 = -0.049
        self.y2 = self.y1
        self.z2 = self.z1

        self.x3 = self.x2
        self.y3 = self.y2
        self.z3 = -0.09

        self.x4 = self.x3
        self.y4 = 0.1305
        self.z4 = -0.21

        self.x5 = -0.0016
        self.y5 = self.y4
        self.z5 = -0.2707

        self.x6 = 0.0184
        self.y6 = self.y5
        self.z6 = -0.7054

        self.x7 = -0.0216
        self.y7 = self.y6
        self.z7 = -1.1134

        self.l8 = 0.05673
        self.x8 = self.x7 + self.l8 * np.sin(0.872509)
        self.y8 = self.y7
        self.z8 = self.z7 - self.l8 * np.cos(0.872509)

        self.w1 = np.array([1.0, 0.0, 0.0])
        self.w2 = np.array([0.0, 0.0, 1.0])
        self.w3 = np.array([0.0, -1.0, 0.0])
        self.w4 = np.array([0.0, -1.0, 0.0])
        self.w5 = np.array([0.0, -1.0, 0.0])
        self.w6 = np.array([0.0, -1.0, 0.0])
        self.w7 = np.array([0.0, -1.0, 0.0])
        self.ws = [self.w1, self.w2, self.w3, self.w4, self.w5, self.w6, self.w7]

        self.p1 = np.array([self.x1, self.y1, self.z1])
        self.p2 = np.array([self.x2, self.y2, self.z2])
        self.p3 = np.array([self.x3, self.y3, self.z3])
        self.p4 = np.array([self.x4, self.y4, self.z4])
        self.p5 = np.array([self.x5, self.y5, self.z5])
        self.p6 = np.array([self.x6, self.y6, self.z6])
        self.p7 = np.array([self.x7, self.y7, self.z7])
        self.p8 = np.array([self.x8, self.y8, self.z8])
        self.ps = [self.p1, self.p2, self.p3, self.p4, self.p5, self.p6, self.p7]

        self.S_list1 = [np.hstack((self.ws[i], np.cross(self.ps[i], self.ws[i]))) for i in range(7)]
        self.S_list2 = [np.hstack((self.ws[i], np.cross(self.ps[i] * [1, -1, 1], self.ws[i]))) for i in range(7)]

        self.S_lists = [self.S_list1, self.S_list2]

        self.M1 = sm.SE3.Trans(self.p8) * sm.SE3.RPY(np.pi / 2, np.pi / 2, 0)
        self.M2 = sm.SE3.Trans(self.p8 * [1, -1, 1]) * sm.SE3.RPY(np.pi / 2, np.pi / 2, 0)
        self.Ms = [self.M1, self.M2]
        self.q60 = 0.2270

        self.dof = 10
        self.q = [0.0 for _ in range(self.dof)]
        self.T = (sm.SE3(), sm.SE3())
        self.set_joint(self.q)
        self.speed = 0.001
        self.mode = 0

    def fkine(self, q_list: list) -> tuple:
        q_list1 = [*(q_list[:4]), 0, self.q60 - q_list[3], q_list[4]]
        q_list2 = [*(q_list[5:9]), 0, self.q60 - q_list[8], q_list[9]]
        return fkine(self.Ms[0], self.S_lists[0], q_list1), fkine(self.Ms[1], self.S_lists[1], q_list2)

    def ikine2(self, S_list, pos, sgn) -> list:

        M1 = sm.SE3.Trans(self.p3 * [1, sgn, 1])
        M2 = sm.SE3.Trans((self.p3 + [0, 0.0045, 0]) * [1, sgn, 1])

        d3 = sgn * (self.y1 - self.y4)

        if self.mode == 0:
            T0 = sm.SE3.Trans(self.ps[0] * [1, sgn, 1])
            pos0 = np.squeeze(T0.inv() * sm.SE3.Trans(pos).t)
            q1 = np.arctan2(-d3, -np.sqrt(np.power(pos0[1], 2) + np.power(pos0[2], 2) - np.power(d3, 2))) \
                 + np.arctan2(-pos0[1], pos0[2])
        else:
            q1 = 0

        q_list = [q1]
        T1 = fkine(M1, S_list[:1], q_list)
        pos1 = np.squeeze(T1.inv() * sm.SE3.Trans(pos).t)

        if self.mode == 0:
            q2 = 0.0
        else:
            q2 = np.arctan2(d3, np.sqrt(np.power(pos1[0], 2) + np.power(pos1[1], 2) - np.power(d3, 2))) \
                 + np.arctan2(pos1[1], pos1[0])

        q_list.append(q2)
        T2 = fkine(M2, S_list[:2], q_list)
        pos2 = np.squeeze(T2.inv() * T1 * sm.SE3.Trans(pos1).t)
        x3 = pos2[0]
        z3 = pos2[2] + self.l8

        l1 = np.linalg.norm((self.p4 - self.p3) * np.array([1.0, 0.0, 1.0]))
        l2 = np.linalg.norm((self.p6 - self.p4) * np.array([1.0, 0.0, 1.0]))
        l3 = np.linalg.norm((self.p7 - self.p6) * np.array([1.0, 0.0, 1.0]))

        q4_0 = np.arctan2((self.p6 - self.p4)[0], -(self.p6 - self.p4)[2])
        q4_2_0 = 0.2270
        q4_2_1 = np.arctan2((self.p7 - self.p6)[0], -(self.p7 - self.p6)[2])

        l0 = np.linalg.norm([x3, z3])
        q0_1 = np.arctan2(x3, -z3)
        q34 = -np.arccos((np.power(l0, 2) + np.power(l2, 2) - np.power(l1, 2) - np.power(l3, 2)
                          - 2 * l1 * l3 * np.cos(q4_2_0 + q4_2_1)) / (2 * l2 * l0)) + q0_1
        k1 = x3 - l2 * np.sin(q34)
        k2 = -z3 - l2 * np.cos(q34)
        l0_2 = np.linalg.norm([k1, k2])
        q0_2 = np.arctan2(k1, k2)
        q3 = -np.arccos((np.power(l0_2, 2) + np.power(l1, 2) - np.power(l3, 2)) / (2 * l1 * l0_2)) + q0_2
        q4 = q34 - q3 - q4_0

        q5_0 = -1.0995
        q5 = q5_0 - q3
        qs = [q1, q2, q3, q4, q5]

        return [wrap_to_pi(q) for q in qs]

    def ikine(self, left, right):
        return [*(self.ikine2(self.S_lists[0], left, 1)), *(self.ikine2(self.S_lists[1], right, -1))]
